package com.example.scribe.backend.handlers;

import com.example.scribe.SyntaxNode;
import com.example.scribe.backend.ASTUtils;
import com.example.scribe.backend.BytecodeOps;
import com.example.scribe.backend.ExpressionEmitter;
import com.example.scribe.backend.ExpressionNodeHandler;
import com.example.scribe.backend.LocalVariableAllocator;
import net.bytebuddy.implementation.bytecode.StackManipulation;
import net.bytebuddy.implementation.bytecode.StackManipulation.AbstractBase;
import net.bytebuddy.implementation.bytecode.StackManipulation.Compound;

/**
 * Handler for comparison and equality expressions (<=, >=, <, >, ==, !=).
 */
public final class ComparisonExpressionHandler implements ExpressionNodeHandler {
    private final ExpressionEmitter exprEmitter;

    public ComparisonExpressionHandler(ExpressionEmitter exprEmitter) {
        this.exprEmitter = exprEmitter;
    }

    @Override
    public boolean canHandle(String nodeType) {
        return "equality".equals(nodeType) || "comparison".equals(nodeType);
    }

    @Override
    public StackManipulation handle(SyntaxNode node, String generatedClassName, LocalVariableAllocator locals) {
        if (node == null || node.children().size() < 3) {
            return StackManipulation.Trivial.INSTANCE;
        }

        SyntaxNode leftNode = node.children().get(0);
        SyntaxNode opNode = node.children().get(1);
        SyntaxNode rightNode = node.children().get(2);
        final String op = opNode.text();

        // Find actual value nodes (identifier or number) under nested wrapper nodes
        StackManipulation left = findValueEmit(leftNode, generatedClassName, locals);
        if (left == StackManipulation.Trivial.INSTANCE) {
            left = BytecodeOps.pushInt(0);
        }
        StackManipulation right = findValueEmit(rightNode, generatedClassName, locals);
        if (right == StackManipulation.Trivial.INSTANCE) {
            right = BytecodeOps.pushInt(0);
        }

        final StackManipulation leftF = left;
        final StackManipulation rightF = right;
        
        return new Compound(leftF, rightF, new AbstractBase() {
            @Override
            public StackManipulation.Size apply(net.bytebuddy.jar.asm.MethodVisitor mv, net.bytebuddy.implementation.Implementation.Context ctx) {
                int opcode;
                switch (op) {
                    case "<=": opcode = net.bytebuddy.jar.asm.Opcodes.IF_ICMPLE; break;
                    case ">=": opcode = net.bytebuddy.jar.asm.Opcodes.IF_ICMPGE; break;
                    case "<": opcode = net.bytebuddy.jar.asm.Opcodes.IF_ICMPLT; break;
                    case ">": opcode = net.bytebuddy.jar.asm.Opcodes.IF_ICMPGT; break;
                    case "==": opcode = net.bytebuddy.jar.asm.Opcodes.IF_ICMPEQ; break;
                    case "!=": opcode = net.bytebuddy.jar.asm.Opcodes.IF_ICMPNE; break;
                    default: opcode = net.bytebuddy.jar.asm.Opcodes.IF_ICMPNE; break;
                }
                net.bytebuddy.jar.asm.Label trueL = new net.bytebuddy.jar.asm.Label();
                net.bytebuddy.jar.asm.Label endL = new net.bytebuddy.jar.asm.Label();
                mv.visitJumpInsn(opcode, trueL);
                mv.visitInsn(net.bytebuddy.jar.asm.Opcodes.ICONST_0);
                mv.visitJumpInsn(net.bytebuddy.jar.asm.Opcodes.GOTO, endL);
                mv.visitLabel(trueL);
                // Frame at jump target: after comparison, stack is empty (both ints consumed by IF_ICMP)
                mv.visitFrame(net.bytebuddy.jar.asm.Opcodes.F_SAME, 0, null, 0, null);
                mv.visitInsn(net.bytebuddy.jar.asm.Opcodes.ICONST_1);
                mv.visitLabel(endL);
                // Frame at goto target: stack has 1 int (comparison result 0 or 1)
                mv.visitFrame(net.bytebuddy.jar.asm.Opcodes.F_SAME1, 0, null, 1, new Object[]{net.bytebuddy.jar.asm.Opcodes.INTEGER});
                return new StackManipulation.Size(1, 0);
            }
        });
    }

    /**
     * Helper method to find and emit value nodes (identifier or number).
     * Similar to the findValueEmit method in BasicExpressionEmitter.
     * This method prioritizes finding literal values (numbers/identifiers) over
     * evaluating complex expressions, which is important for comparison operations.
     */
    private StackManipulation findValueEmit(SyntaxNode node, String generatedClassName, LocalVariableAllocator locals) {
        if (node == null) {
            return StackManipulation.Trivial.INSTANCE;
        }

        // Direct check for leaf nodes first (prioritize literals)
        if (ASTUtils.isNumber(node)) {
            try {
                int v = Integer.parseInt(node.text());
                return BytecodeOps.pushInt(v);
            } catch (NumberFormatException e) {
                return StackManipulation.Trivial.INSTANCE;
            }
        }
        
        if (ASTUtils.isIdentifier(node)) {
            String name = node.text();
            if (locals.hasName(name)) {
                return BytecodeOps.loadInt(locals.getIndexByName(name));
            }
        }

        // Check Call nodes FIRST - they may wrap primary nodes with numbers/identifiers
        // This is important for literal values like 1, 0, etc.
        // Structure: Call -> primary -> Terminal:NUMBER
        // Priority: Check Call before other structured nodes to ensure literals are found
        if ("Call".equals(node.type())) {
            for (SyntaxNode child : node.children()) {
                // First check if child is primary node and directly contains number/identifier
                if ("primary".equals(child.type())) {
                    for (SyntaxNode prim : child.children()) {
                        if (ASTUtils.isNumber(prim)) {
                            try {
                                int v = Integer.parseInt(prim.text());
                                return BytecodeOps.pushInt(v);
                            } catch (NumberFormatException e) {
                                // continue
                            }
                        }
                        if (ASTUtils.isIdentifier(prim)) {
                            String name = prim.text();
                            if (name != null && locals.hasName(name)) {
                                return BytecodeOps.loadInt(locals.getIndexByName(name));
                            }
                        }
                    }
                }
                // Always recursively search all Call children (including primary)
                // This ensures we find nested literals even if they're deeper
                StackManipulation nested = findValueEmit(child, generatedClassName, locals);
                if (nested != StackManipulation.Trivial.INSTANCE) {
                    return nested;
                }
            }
        }

        // Check primary nodes for identifier or number
        if ("primary".equals(node.type())) {
            for (SyntaxNode child : node.children()) {
                if (ASTUtils.isNumber(child)) {
                    try {
                        int v = Integer.parseInt(child.text());
                        return BytecodeOps.pushInt(v);
                    } catch (NumberFormatException e) {
                        // continue
                    }
                }
                if (ASTUtils.isIdentifier(child)) {
                    String name = child.text();
                    if (name != null && locals.hasName(name)) {
                        return BytecodeOps.loadInt(locals.getIndexByName(name));
                    }
                }
            }
            // Also recursively search primary children
            for (SyntaxNode child : node.children()) {
                StackManipulation nested = findValueEmit(child, generatedClassName, locals);
                if (nested != StackManipulation.Trivial.INSTANCE) {
                    return nested;
                }
            }
        }

        // For structured nodes (comparison, equality, etc.), use emit directly
        if ("comparison".equals(node.type()) || "equality".equals(node.type())) {
            StackManipulation result = exprEmitter.emit(node, generatedClassName, locals);
            if (result != StackManipulation.Trivial.INSTANCE) {
                return result;
            }
        }

        // For expression nodes, first check if it contains comparison/equality
        if ("expression".equals(node.type())) {
            for (SyntaxNode child : node.children()) {
                if ("comparison".equals(child.type()) || "equality".equals(child.type())) {
                    StackManipulation result = exprEmitter.emit(child, generatedClassName, locals);
                    if (result != StackManipulation.Trivial.INSTANCE) {
                        return result;
                    }
                }
            }
        }

        // For structured nodes like shift, multiplication, unary, etc., recursively search children
        // These nodes may wrap the actual literal values
        // Note: unary nodes may wrap Call nodes which wrap primary nodes
        // Important: bitwise_or, bitwise_xor, bitwise_and may have only one child (no operator)
        // In that case, we need to recursively search to find the actual literal
        if ("shift".equals(node.type()) || "multiplication".equals(node.type()) || 
            "unary".equals(node.type()) || "addition".equals(node.type()) ||
            "subtraction".equals(node.type()) || "division".equals(node.type()) ||
            "bitwise_or".equals(node.type()) || "bitwise_xor".equals(node.type()) ||
            "bitwise_and".equals(node.type())) {
            // For these nodes, if they have only one child (no operator), 
            // recursively search to find the literal value
            if (node.children().size() == 1) {
                // Only one child - no operator, just pass through
                return findValueEmit(node.children().get(0), generatedClassName, locals);
            } else {
                // Multiple children - has operator, recursively search all children
                // (though for comparison, we only want the first operand)
                for (SyntaxNode child : node.children()) {
                    // Skip operator terminals
                    if (!child.type().startsWith("Terminal:")) {
                        StackManipulation nested = findValueEmit(child, generatedClassName, locals);
                        if (nested != StackManipulation.Trivial.INSTANCE) {
                            return nested;
                        }
                    }
                }
            }
        }

        // Recursively search all children for literals before trying emit
        for (SyntaxNode child : node.children()) {
            // Skip operator terminals in recursive search
            if (!child.type().startsWith("Terminal:")) {
                StackManipulation nested = findValueEmit(child, generatedClassName, locals);
                if (nested != StackManipulation.Trivial.INSTANCE) {
                    return nested;
                }
            }
        }

        // Try direct emit for other structured nodes as last resort
        // But skip emit for shift/multiplication/unary as they may not have operators
        if (!"shift".equals(node.type()) && !"multiplication".equals(node.type()) && 
            !"unary".equals(node.type()) && !"addition".equals(node.type()) &&
            !"subtraction".equals(node.type()) && !"division".equals(node.type())) {
            StackManipulation result = exprEmitter.emit(node, generatedClassName, locals);
            if (result != StackManipulation.Trivial.INSTANCE) {
                return result;
            }
        }
        
        return StackManipulation.Trivial.INSTANCE;
    }
}

